package Client;

import java.math.BigInteger;
import java.security.SecureRandom;
import org.bouncycastle.crypto.BlockCipher;
import org.bouncycastle.crypto.BufferedAsymmetricBlockCipher;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SHA1Digest;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.engines.RSAEngine;
import org.bouncycastle.crypto.generators.PKCS5S2ParametersGenerator;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PKCS7Padding;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.crypto.params.RSAKeyParameters;

/**
 * Implementation of cryptographic functions
 * needed to communicate with the Smart Dating Service
 * @author nvamelichev
 */
public class CryptoImpl {
  /**
   * @return cryptographically random session key for exchange
   * between client and service
   */
  static byte[] getSessionKey() {
    byte[] sessionKey = new byte[Constants.KEY_SIZE / 8];
    SecureRandom random = new SecureRandom();
    random.nextBytes(sessionKey);
    return sessionKey;
  }

  static byte[] cryptAes(byte[] data, byte[] sessionKey) throws DataLengthException, IllegalStateException, InvalidCipherTextException {
    PKCS5S2ParametersGenerator generator = new PKCS5S2ParametersGenerator();
    generator.init(sessionKey, Constants.SALT, Constants.ITERATIONS);
    CipherParameters generatedParams =
            generator.generateDerivedParameters(Constants.KEY_SIZE);
    
    BlockCipher engine = new AESEngine();
    PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(engine), new PKCS7Padding());
    ParametersWithIV params = new ParametersWithIV(generatedParams, Constants.ENCRYPTION_IV);

    cipher.init(true, params);
    
      byte[] out = new byte[cipher.getOutputSize(data.length)];
      int len = cipher.processBytes(data, 0, data.length, out, 0);
      len += cipher.doFinal(out, len);

      // remove excessive padding
      byte[] result = new byte[len];
      System.arraycopy(out, 0, result, 0, len);
      return result;
  }

  static byte[] decryptAes(byte[] encryptedData, byte[] sessionKey) {
    PKCS5S2ParametersGenerator generator = new PKCS5S2ParametersGenerator();
    generator.init(sessionKey, Constants.SALT, Constants.ITERATIONS);
    CipherParameters generatedParams =
            generator.generateDerivedParameters(Constants.KEY_SIZE);

    BlockCipher engine = new AESEngine();
    PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(engine), new PKCS7Padding());
    ParametersWithIV params = new ParametersWithIV(generatedParams, Constants.ENCRYPTION_IV);

    cipher.init(false, params);

    byte[] out = new byte[cipher.getOutputSize(encryptedData.length)];
    int len = cipher.processBytes(encryptedData, 0, encryptedData.length, out, 0);

    // remove excessive padding
    byte[] result = new byte[len];
    System.arraycopy(out, 0, result, 0, len);
    return result;
  }

  static byte[] sha1(byte[] data) {
    SHA1Digest digest = new SHA1Digest();

    digest.update(data, 0, data.length);
    byte[] result = new byte[digest.getDigestSize()];
    digest.doFinal(result, 0);
    
    return result;
  }

  // Taken from openjdk, see RSAPadding#padV15(byte[]) for reference
  // http://www.docjar.com/html/api/sun/security/rsa/RSAPadding.java.html
  private static byte[] padPKCS15(byte[] data, int type) {
    int paddedSize = Constants.SERVER_PUBLIC_MODULUS.length;
    byte[] padded = new byte[paddedSize];
    System.arraycopy(data, 0, padded, paddedSize - data.length, data.length);
    int psSize = paddedSize - 3 - data.length;
    int k = 0;
    padded[k++] = 0;
    padded[k++] = (byte)type;
    if (type == 1) {
     // blocktype 1: all padding bytes are 0xff
     while (psSize-- > 0) {
         padded[k++] = (byte)0xff;
     }
    } else {
     // blocktype 2: padding bytes are random non-zero bytes
     SecureRandom random = new SecureRandom();
     // generate non-zero padding bytes
     // use a buffer to reduce calls to SecureRandom
     byte[] r = new byte[64];
     int i = -1;
     while (psSize-- > 0) {
         int b;
         do {
             if (i < 0) {
                 random.nextBytes(r);
                 i = r.length - 1;
             }
             b = r[i--] & 0xff;
         } while (b == 0);
         padded[k++] = (byte)b;
     }
    }
    return padded;
  }

  static byte[] rsaForKeyExchange(byte[] sessionKey) throws InvalidCipherTextException {
    RSAEngine engine = new RSAEngine();
    BufferedAsymmetricBlockCipher cipher = new BufferedAsymmetricBlockCipher(engine);

    byte[] data = padPKCS15(sessionKey, 2);
    cipher.init(true, new RSAKeyParameters(false, new BigInteger(1,Constants.SERVER_PUBLIC_MODULUS),
            Constants.SERVER_PUBLIC_EXPONENT));
    cipher.processBytes(data, 0, data.length);
    return cipher.doFinal();
  }
}
